import numpy as np
import tensorflow.keras.backend as K
from tensorflow.keras.layers import (Activation, BatchNormalization, Conv2D,
                                     Input,
                                     Reshape, Concatenate, ELU, MaxPooling2D)
from tensorflow.keras.models import Model


def conv2d(inputs,
           filters=32,
           kernel_size=3,
           strides=1,
           name=None):
    conv = Conv2D(filters=filters,
                  kernel_size=kernel_size,
                  strides=strides,
                  kernel_initializer='he_normal',
                  name=name,
                  padding='same')

    return conv(inputs)


def conv_layer(inputs,
               filters=32,
               kernel_size=3,
               strides=1,
               use_maxpool=True,
               postfix=None,
               activation=None):
    x = conv2d(inputs,
               filters=filters,
               kernel_size=kernel_size,
               strides=strides,
               name='conv' + postfix)
    x = BatchNormalization(name="bn" + postfix)(x)
    x = ELU(name='elu' + postfix)(x)
    if use_maxpool:
        x = MaxPooling2D(name='pool' + postfix)(x)
    return x


def build_ssd(input_shape, backbone, n_layers, n_classes,
              aspect_ratios=(1, 2, 0.5)):
    n_anchors = len(aspect_ratios) + 1

    inputs = Input(input_shape)
    outputs = []
    feature_shapes = []
    out_cls = []
    out_off = []

    base_outputs = backbone(inputs)

    for i in range(n_layers):
        conv = base_outputs if n_layers == 1 else base_outputs[i]
        name = "cls" + str(i + 1)
        classes = conv2d(conv, n_anchors * n_classes, kernel_size=3, name=name)

        name = "off" + str(i + 1)
        offsets = conv2d(conv, n_anchors * 4, kernel_size=3, name=name)

        shape = np.array(K.int_shape(offsets))[1:]
        feature_shapes.append(shape)

        name = "cls_res" + str(i + 1)
        classes = Reshape((-1, n_classes), name=name)(classes)

        name = "off_res" + str(i + 1)
        offsets = Reshape((-1, 4), name=name)(offsets)

        offsets = [offsets, offsets]
        name = "off_cat" + str(i + 1)
        offsets = Concatenate(axis=-1, name=name)(offsets)

        out_off.append(offsets)

        name = "cls_out" + str(i + 1)

        classes = Activation('softmax', name=name)(classes)

        out_cls.append(classes)

    if n_layers > 1:
        name = "offsets"
        offsets = Concatenate(axis=1, name=name)(out_off)
        name = "classes"
        classes = Concatenate(axis=1, name=name)(out_cls)

    else:
        offsets = out_off[0]
        classes = out_cls[0]

    outputs = [classes, offsets]

    model = Model(inputs=inputs, outputs=outputs, name='sdd_head')

    return n_anchors, feature_shapes, model
